@wsmoses could there be a pass in :all mode that reverts the performance gain in the passes of :before_enzyme and :after_enzyme?

this benchmark is mostly only dot_general and involved transposes.

Yeah potentially, definitely merits investigation.

cc @avik-pal esp cuz fft related issue is similar

Here is the MLIR generated for each case: vqe-mlir.zip

I'm running into problems with xprof. I'm doing the following:

Reactant.with_profiler(@__DIR__; trace_host=true, trace_device=true) do
    ∇f_xla = @compile compile_options = Reactant.DefaultXLACompileOptions(; sync=true) ∇expectation(
        params_re, observable_re, coef_re
    )
    for _ in 1:100
        ∇f_xla(params_re, observable_re, coef_re)
    end
end

and I get the following message and no data on my traces:

The gradients looks quite strange. Did some loop or vector op get scalarized?

The gradients looks quite strange. Did some loop or vector op get scalarized?

Ahh that's due to the fact that I was zero initializing the parameters and just using one Hamiltonian term; i.e. the ket and the bra are then effectively perpendicular and thus, the primal value and the gradients are zero.

Random initialization shows some better but small gradients. Most probably I would need to add more Hamiltonian terms but right now, we are just sequentially running the gradient function over all the Hamiltonian terms (some parallelization using MPI) and then summing. Batching over the Hamiltonian terms requires some more work but this benchmark (which you can imagine as just running runs 1 epoch, 1 sample) is a good reflection of what we do.

tldr: Gradients being zero are / were a numerical issue, not a bug introduced by a pass.

There is a bug introduced in EnzymeAD/Enzyme-JAX#1121 breaks the rev diff rule of stablehlo.dynamic_update_slice and thus, this benchmark is currently fixed to Reactant v0.2.138

For 128 hamiltonian terms, 40 qubits and 6 layers of the EfficientSU2 ansatz, I'm getting the following results:

The same for 50 qubits seems to have problems:

It could also be noise as I've tested in my laptop and on CPU, and as the benchmarks take longer, less samples are taken.